// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef CLUSTER_DISPATCHER_DISPATCHER_H_
#define CLUSTER_DISPATCHER_DISPATCHER_H_

#include <string>
#include <unordered_map>
#include <memory>
#include <fstream>
#include "cluster/basic/ai_sim.h"
#include "cluster/basic/topology.h"
#include "cluster/network_controller/topo_manager/topo_manager.h"
#include "cluster/resource_scheduler/resource_scheduler.h"
#include "cluster/task_scheduler/task_allocator.h"
#include "cluster/network_controller/network_routing.h"
#include "cluster/task_scheduler/task_type/allreduce.h"
#include "cluster/network_controller/routing_algorithm/route.h"
#include "cluster/roce/RoceMsg_m.h"
#include "flowparser.h"
// #include "cluster/basic/http_get.h"

namespace ai_sim {

class ResourceScheduler;
class TaskAllocator;
class NetworkRouting;

class Dispatcher : public omnetpp::cSimpleModule {
 private:
    ResourceScheduler* resourceScheduler;
    TaskAllocator* taskAllocator;
    NetworkRouting* networkRouting;
    TopoManager* topos;
    ClusterType clustertype;
    std::unordered_map<int, taskboard> taskboards;
/*    
   0 = gen tasks by task gen
   1 = get flows from txt
   2 = get tasks from http server
 */
    int inputType;

    // For inputType = 1
    FlowParser* flowparser;
    ifstream routefile;

 public:
    void initialize() override;
    void handleMessage(omnetpp::cMessage *msg) override;

    // Calls allocFlows() to alloc all flows of the task
    void allocAllTaskFlows(taskboard&);
    // Alloc specific flows to target cpu
    void allocFlows(map<int, FlowVector> tx_flows,
      map<int, FlowVector> rx_flows);

    // Core function of the Dispatcher
    // Alloc flows from file
    void alloc_flow_from_txt();
    // Alloc tasks from alloc gpu tp alloc flow and routes
    void alloc_task(const Task&);
    // Get the new task from the http server(maybe python)
   //  void alloc_task_from_http();

    // Release the resources that a task used after the task
    // finished
    void try_release_task(const Task&);
    void get_statistics(const taskboard&);
};
}  // namespace ai_sim

#endif  // CLUSTER_DISPATCHER_DISPATCHER_H_
